﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Digraphs;
using Algoritmia.DataStructures.Maps;
using Algoritmia.Problems.ConnectedComponents;
using Algoritmia.Utils;

namespace Algoritmia.Problems.SpanningTrees.Baruvka
{
    public class BaruvkasMinimumSpanningForestFinder<T> : IMinimumSpanningForestFinder<T,double>
    {
        private readonly Func<IDigraph<T>, IConnectedComponentsFinder<T>> _createConnectedComponentsFinder;
        private readonly Func<IEnumerable<T>, IEditableDigraph<T>> _createEditableUndirectedGraph;
        private readonly Func<IEnumerable<T>, IMap<T, int>> _createMap;


        public BaruvkasMinimumSpanningForestFinder(Func<IDigraph<T>,IConnectedComponentsFinder<T>> createConnectedComponentsFinder = null,
                                                   Func<IEnumerable<T>,IEditableDigraph<T>> createEditableUndirectedGraph = null,
                                                   Func<IEnumerable<T>,IMap<T,int>> createMap = null)
        {
            _createConnectedComponentsFinder = createConnectedComponentsFinder ??
                                               (g => new GraphTraversalConnectedComponentsFinder<T>());
            _createEditableUndirectedGraph = createEditableUndirectedGraph ??
                (it => new UndirectedGraph<T>(it,Enumerable.Empty<Tuple<T,T>>()));
            _createMap = createMap ?? (it => AlgoritmiaFactory.Maps<T, int>.Default);
        }

        #region IMinimumSpanningForestFinder<T,double> Members

        public IEnumerable<Tuple<T,T>> MinimumSpanningForest(IDigraph<T> g, WeightingFunction<T, double> w)
        {
            var g2 = _createEditableUndirectedGraph(g.V);
            var connectedComponentsFinder = _createConnectedComponentsFinder(g);
            bool[] merged;
            do
            {
                var components = connectedComponentsFinder.ConnectedComponents(g2).ToList();
                var lbl = _createMap(g.V);
                for (int i = 0; i < components.Count; i++)
                {
                    foreach (var vertex in components[i])
                    {
                        lbl[vertex] = i;
                    }
                }
                merged = new bool[components.Count];
                for (int i = 0; i < components.Count; i++)
                {
                    if (!merged[i])
                    {
                        bool noelement = false;
                        Tuple<T, T> e = default(Tuple<T,T>);
                        try
                        {
                            e = (from u in components[i]
                                     from v in g.Successors(u)
                                     where lbl[u] != lbl[v]
                                     select new Tuple<T, T>(u, v)).ArgMin(t => w(t.Item1, t.Item2));
                        }
                        catch
                        {
                            noelement = true;
                        }
                        if (!noelement)
                        {
                            g2.E.Add(e);
                            yield return e;
                            merged[lbl[e.Item2]] = true;
                            foreach (var v in components[lbl[e.Item2]])
                            {
                                lbl[v] = lbl[e.Item1];
                            }
                        }
                    }
                }

            } while (merged.Any(e => e));
        }

        #endregion
    }
}
